#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gravity calibration & online compensation (ROS2, rclpy)
同步补偿：/wrench_comp 发布频率与 /force_data 保持一致
"""
import os
import json
from pathlib import Path
import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.duration import Duration
from rclpy.time import Time

import tf2_ros
from geometry_msgs.msg import WrenchStamped
from std_srvs.srv import Trigger


def quat_to_rotmat(x, y, z, w):
    n = np.sqrt(x * x + y * y + z * z + w * w)
    if n == 0.0:
        return np.eye(3, dtype=float)
    x, y, z, w = x / n, y / n, z / n, w / n
    xx, yy, zz = x * x, y * y, z * z
    xy, xz, yz = x * y, x * z, y * z
    wx, wy, wz = w * x, w * y, w * z
    return np.array([
        [1 - 2 * (yy + zz), 2 * (xy - wz),     2 * (xz + wy)],
        [2 * (xy + wz),     1 - 2 * (xx + zz), 2 * (yz - wx)],
        [2 * (xz - wy),     2 * (yz + wx),     1 - 2 * (xx + yy)]
    ], dtype=float)


CALIB_FILE = Path.home() / ".ros" / "gravity_calib.json"


class GravityCalibNode(Node):
    def __init__(self):
        super().__init__('gravity_calib_node')

        # ---------------- 参数 ----------------
        self.declare_parameter('base_frame', 'base_link')
        self.declare_parameter('ee_frame', 'end')
        self.declare_parameter('wrench_topic', '/force_data')
        self.declare_parameter('capture_timeout_s', 2.0)
        self.declare_parameter('g_B', None)
        self.declare_parameter('F0', None)

        self.declare_parameter('enable_online_bias_tune', True)
        self.declare_parameter('no_contact_thresh', 1.0)
        self.declare_parameter('alpha_bias', 0.001)

        self.base_frame = self.get_parameter('base_frame').value
        self.ee_frame = self.get_parameter('ee_frame').value
        self.wrench_topic = self.get_parameter('wrench_topic').value
        self.cap_timeout = float(self.get_parameter('capture_timeout_s').value)

        self.enable_online_bias_tune = bool(self.get_parameter('enable_online_bias_tune').value)
        self.no_contact_thresh = float(self.get_parameter('no_contact_thresh').value)
        self.alpha_bias = float(self.get_parameter('alpha_bias').value)

        # 末端->力传感器固定旋转
        self.M = np.array([[0, 0, 1],
                           [0, -1, 0],
                           [1, 0, 0]], dtype=float)

        # ---------------- TF / Topics ----------------
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)

        self.last_force = None
        self.sub_wrench = self.create_subscription(
            WrenchStamped, self.wrench_topic, self.cb_wrench, 10
        )
        self.pub_comp = self.create_publisher(WrenchStamped, '/wrench_comp', 10)

        # ---------------- 标定与样本 ----------------
        self.samples_RSB = []
        self.samples_F = []

        self.g_B = None
        self.F0 = None

        self._load_calib_from_params_or_file()

        # ---------------- 服务 ----------------
        self.srv_capture = self.create_service(Trigger, 'capture_sample', self.handle_capture)
        self.srv_solve = self.create_service(Trigger, 'solve_calibration', self.handle_solve)
        self.srv_save = self.create_service(Trigger, 'save_calibration', self.handle_save)
        self.srv_clear = self.create_service(Trigger, 'clear_calibration', self.handle_clear)
        self.srv_reload = self.create_service(Trigger, 'reload_calibration', self.handle_reload)

        self.get_logger().info(
            f'GravityCalibNode up. base={self.base_frame}, ee={self.ee_frame}, '
            f'wrench={self.wrench_topic}, calib_file={str(CALIB_FILE)}'
        )

    # ===== 基础回调 =====
    def cb_wrench(self, msg: WrenchStamped):
        self.last_force = np.array([msg.wrench.force.x,
                                    msg.wrench.force.y,
                                    msg.wrench.force.z], dtype=float)
        # 收到新数据时立即做补偿并发布
        self.online_compensation(msg)

    def lookup_R_SB(self) -> np.ndarray:
        now = Time()
        ok = self.tf_buffer.can_transform(self.base_frame, self.ee_frame, now,
                                          timeout=Duration(seconds=self.cap_timeout))
        if not ok:
            raise RuntimeError(f'等待 TF 超时: {self.base_frame}->{self.ee_frame}')

        trans = self.tf_buffer.lookup_transform(self.base_frame, self.ee_frame, now,
                                                timeout=Duration(seconds=0.5))
        q = trans.transform.rotation
        R_BE = quat_to_rotmat(q.x, q.y, q.z, q.w)
        R_BS = R_BE @ self.M
        R_SB = R_BS.T
        return R_SB

    # ===== 标定采样/求解 =====
    def handle_capture(self, request, response):
        try:
            if self.last_force is None:
                raise RuntimeError('暂未收到力传感器数据')
            R_SB = self.lookup_R_SB()
            self.samples_RSB.append(R_SB)
            self.samples_F.append(self.last_force.copy())
            n = len(self.samples_F)
            response.success = True
            response.message = f'采样成功，第 {n} 组。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'采样失败: {e}'
            self.get_logger().warn(response.message)
        return response

    def handle_solve(self, request, response):
        try:
            N = len(self.samples_F)
            if N < 3:
                raise RuntimeError(f'样本不足，需要≥3，当前 {N}')

            A_rows, b_rows = [], []
            for i in range(N):
                R = self.samples_RSB[i]
                F = self.samples_F[i]
                A_rows.append(np.hstack([R, np.eye(3)]))
                b_rows.append(F)
            A = np.vstack(A_rows)
            b = np.hstack(b_rows)

            x, *_ = np.linalg.lstsq(A, b, rcond=None)
            self.g_B = x[0:3]
            self.F0 = x[3:6]

            g_norm = float(np.linalg.norm(self.g_B))
            self.get_logger().info(f'求解完成: g_B={self.g_B} (|g_B|={g_norm:.4f}), F0={self.F0}')

            self._save_calib_to_file()

            response.success = True
            response.message = '标定完成并已保存，进入在线补偿模式。'
        except Exception as e:
            response.success = False
            response.message = f'求解失败: {e}'
            self.get_logger().error(response.message)
        return response

    # ===== 文件持久化与维护服务 =====
    def _validate_vec(self, v):
        return isinstance(v, (list, tuple)) and len(v) == 3

    def _load_calib_from_params_or_file(self):
        g_param = self.get_parameter('g_B').value
        F_param = self.get_parameter('F0').value
        if self._validate_vec(g_param) and self._validate_vec(F_param):
            self.g_B = np.array(g_param, dtype=float)
            self.F0 = np.array(F_param, dtype=float)
            self.get_logger().info(f'从参数加载标定：g_B={self.g_B}, F0={self.F0}')
            return

        if CALIB_FILE.exists():
            try:
                with open(CALIB_FILE, 'r') as f:
                    data = json.load(f)
                if (self._validate_vec(data.get('g_B')) and
                        self._validate_vec(data.get('F0')) and
                        data.get('base_frame') == self.base_frame and
                        data.get('ee_frame') == self.ee_frame):
                    self.g_B = np.array(data['g_B'], dtype=float)
                    self.F0 = np.array(data['F0'], dtype=float)
                    self.get_logger().info(f'从文件加载标定：g_B={self.g_B}, F0={self.F0}')
                else:
                    self.get_logger().warn('标定文件存在，但与当前配置不匹配或字段缺失，忽略。')
            except Exception as e:
                self.get_logger().warn(f'读取标定文件失败：{e}')

    def _save_calib_to_file(self):
        if self.g_B is None or self.F0 is None:
            raise RuntimeError('当前无标定数据可保存')
        CALIB_FILE.parent.mkdir(parents=True, exist_ok=True)
        data = {
            'g_B': self.g_B.tolist(),
            'F0': self.F0.tolist(),
            'base_frame': self.base_frame,
            'ee_frame': self.ee_frame,
            'M': self.M.tolist()
        }
        with open(CALIB_FILE, 'w') as f:
            json.dump(data, f, indent=2)
        self.get_logger().info(f'标定已保存：{CALIB_FILE}')

    def handle_save(self, request, response):
        try:
            self._save_calib_to_file()
            response.success = True
            response.message = '已保存到文件。'
        except Exception as e:
            response.success = False
            response.message = f'保存失败: {e}'
        return response

    def handle_clear(self, request, response):
        try:
            if CALIB_FILE.exists():
                CALIB_FILE.unlink()
            self.g_B, self.F0 = None, None
            response.success = True
            response.message = '已清除本地标定与内存。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'清除失败: {e}'
            self.get_logger().error(response.message)
        return response

    def handle_reload(self, request, response):
        try:
            self.g_B, self.F0 = None, None
            self._load_calib_from_params_or_file()
            ok = (self.g_B is not None and self.F0 is not None)
            response.success = ok
            response.message = '已重载标定。' if ok else '未找到可用标定。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'重载失败: {e}'
            self.get_logger().error(response.message)
        return response

    # ===== 在线补偿（订阅触发）=====
    def online_compensation(self, msg_in: WrenchStamped):
        if self.g_B is None or self.F0 is None:
            return
        try:
            R_SB = self.lookup_R_SB()
            F_pred = R_SB @ self.g_B + self.F0
            F_comp = self.last_force - F_pred

            if self.enable_online_bias_tune:
                if np.linalg.norm(F_comp) < self.no_contact_thresh:
                    self.F0 = (1.0 - self.alpha_bias) * self.F0 + \
                              self.alpha_bias * (self.last_force - R_SB @ self.g_B)

            msg_out = WrenchStamped()
            msg_out.header.stamp = msg_in.header.stamp  # 使用 force_data 时间戳
            msg_out.header.frame_id = 'sensor'
            msg_out.wrench.force.x = -float(F_comp[0])
            msg_out.wrench.force.y = float(F_comp[1])
            msg_out.wrench.force.z = -float(F_comp[2])
            self.pub_comp.publish(msg_out)

            self.get_logger().info(f"wrench_comp:: Fx={msg_out.wrench.force.x:.3f}, Fz={msg_out.wrench.force.z:.3f}")
        except Exception as e:
            self.get_logger().warn(f'在线补偿失败: {e}')


def main():
    rclpy.init()
    node = GravityCalibNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
